package com.inzone.fram.gray.common.filter;

import com.inzone.fram.gray.common.rule.MatchingContext;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.LoadBalancerClient;
import org.springframework.cloud.gateway.config.LoadBalancerProperties;
import org.springframework.cloud.gateway.filter.LoadBalancerClientFilter;
import org.springframework.http.HttpHeaders;
import org.springframework.web.server.ServerWebExchange;

import java.net.URI;
import java.util.List;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.GATEWAY_REQUEST_URL_ATTR;

/**
 * 负载均衡拦截器
 * 需要重写网关原来的负载均衡拦截器
 * 否则无法使用基于ThreadLocal的请求上下文
 * @author 丁常磊
 * @date 2021/8/23 18:56
 */
public class GrayLoadBalanceFilter extends LoadBalancerClientFilter {

    public GrayLoadBalanceFilter(LoadBalancerClient loadBalancer, LoadBalancerProperties properties) {
        super(loadBalancer, properties);
    }

    @Override
    protected ServiceInstance choose(ServerWebExchange exchange) {

        //使用当前请求头初始化命中上下文
        MatchingContext context = MatchingContext.init();

        //获取当前请求头
        HttpHeaders headers = exchange.getRequest().getHeaders();

        List<String> list = headers.get(MatchingContext.GRAY_HEADER_GROUPS);

        if (list != null && list.size() > 0) {
            String s = list.get(0);
            context.setGroupsText(s);
        }

        list = headers.get(MatchingContext.GRAY_HEADER_TAGS);

        if (list != null && list.size() > 0) {
            String s = list.get(0);
            context.setTagsText(s);
        }

        try {
            return loadBalancer.choose(
                    ((URI) exchange.getAttribute(GATEWAY_REQUEST_URL_ATTR)).getHost());
        }finally {
            //重置命中上下文
            MatchingContext.reset();
        }
    }
}
